import json
import os
import argparse
import logging
from tqdm import tqdm  # 用于显示进度条

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 标签映射（根据实际标注修改）
name2id = {'0': 0, '1': 1, '2': 2}  # 更新为0,1,2三个标签


def polygon_to_bbox(points):
    """计算多边形的最小外接矩形（x1,y1,x2,y2）"""
    xs = [p[0] for p in points]
    ys = [p[1] for p in points]
    return (min(xs), min(ys), max(xs), max(ys))  # 左上、右下坐标


def rectangle_to_bbox(points):
    """将矩形标注（两个点）转换为边界框（x1,y1,x2,y2）"""
    x1, y1 = points[0]  # 左上角点
    x2, y2 = points[1]  # 右下角点
    return (min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2))


def convert_to_yolo(img_size, bbox):
    """将包围盒转换为YOLO归一化坐标（x_center, y_center, width, height）"""
    x1, y1, x2, y2 = bbox
    dw, dh = 1.0 / img_size[0], 1.0 / img_size[1]
    x = (x1 + x2) / 2.0 * dw
    y = (y1 + y2) / 2.0 * dh
    w = (x2 - x1) * dw
    h = (y2 - y1) * dh
    return (x, y, w, h)


def process_json(json_path, output_dir):
    """处理单个JSON文件，生成YOLO标注"""
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except Exception as e:
        logger.error(f"无法解析JSON文件 {json_path}: {str(e)}")
        return False

    img_w, img_h = data.get('imageWidth'), data.get('imageHeight')
    if not img_w or not img_h:
        logger.error(f"JSON文件 {json_path} 缺少图像尺寸信息")
        return False

    txt_path = os.path.join(output_dir, os.path.basename(json_path)[:-5] + '.txt')
    valid_annotations = False

    try:
        with open(txt_path, 'w') as txt_file:
            for shape in data.get('shapes', []):
                shape_type = shape.get('shape_type')
                label = shape.get('label')

                # 检查标签是否存在
                if not label or label not in name2id:
                    logger.warning(f"跳过未知标签 '{label}' in {json_path}")
                    continue

                # 根据形状类型处理标注
                if shape_type == 'polygon':
                    bbox = polygon_to_bbox(shape.get('points', []))
                elif shape_type == 'rectangle':
                    bbox = rectangle_to_bbox(shape.get('points', []))
                else:
                    logger.warning(f"跳过不支持的形状类型 '{shape_type}' in {json_path}")
                    continue

                yolo_coords = convert_to_yolo((img_w, img_h), bbox)
                txt_file.write(f"{name2id[label]} {' '.join(f'{v:.6f}' for v in yolo_coords)}\n")
                valid_annotations = True

        # 清理空文件（无有效标注时）
        if not valid_annotations and os.path.exists(txt_path):
            os.remove(txt_path)
            logger.info(f"删除空文件：{txt_path}")
            return False

        return True

    except Exception as e:
        logger.error(f"处理文件 {json_path} 时出错: {str(e)}")
        return False


def main():
    parser = argparse.ArgumentParser(description='将LabelMe JSON标注转换为YOLO格式')
    parser.add_argument('--json_dir', required=True, help='LabelMe JSON文件目录')
    parser.add_argument('--output_dir', required=True, help='YOLO标注输出目录')
    parser.add_argument('--labels', default='0,1,2',
                        help='逗号分隔的标签列表，例如: "0,1,2"')
    args = parser.parse_args()

    # 更新标签映射
    global name2id
    labels = args.labels.split(',')
    name2id = {label: i for i, label in enumerate(labels)}
    logger.info(f"使用标签映射: {name2id}")

    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)

    # 获取所有JSON文件
    json_files = [f for f in os.listdir(args.json_dir) if f.endswith('.json')]
    total_files = len(json_files)

    if total_files == 0:
        logger.warning(f"在目录 {args.json_dir} 中未找到JSON文件")
        return

    # 处理所有JSON文件
    success_count = 0
    failed_files = []

    logger.info(f"开始处理 {total_files} 个JSON文件...")

    for json_file in tqdm(json_files, desc="转换进度"):
        json_path = os.path.join(args.json_dir, json_file)
        if process_json(json_path, args.output_dir):
            success_count += 1
        else:
            failed_files.append(json_file)

    # 输出统计结果
    logger.info(f"转换完成！成功: {success_count}, 失败: {len(failed_files)}, 总计: {total_files}")

    if failed_files:
        logger.warning(f"失败的文件 ({len(failed_files)}):")
        for file in failed_files[:10]:  # 只显示前10个失败文件
            logger.warning(f"  - {file}")
        if len(failed_files) > 10:
            logger.warning(f"  - ... 等 {len(failed_files) - 10} 个文件")


if __name__ == "__main__":
    main()